#include "config.h"

#include <clx/libECM.h>
#include <gsl/gsl_fit.h>
#include <math.h>

#define ERR(fmt, ...)	ECMERR("ERR: lstatpop: " fmt, ## __VA_ARGS__)

static char	*inputfile;

static ecm_uint32	n_inputs;

static ecm_double	*ranks;
static ecm_double	*accesses;

static void
usage(void)
{
	ecm_fprintf(ecm_stderr,
"Usage: lstatpop [ <options> ] <arrival file>\n"
" Options:\n"
"   -h: help\n"
"   -o: output file\n"
	);
}

#define ALLOC_CHUNK	256

static ECMBOOL
read_input(void)
{
	FILE	*fp;
	char	buf[256];
	ecm_uint32	n_alloced = ALLOC_CHUNK;
	ecm_uint32	i;

	fp = ecm_fopen(inputfile, "r");
	if (fp == NULL) {
		ERR("%s: cannot open file", inputfile);
		return ECM_FALSE;
	}

	accesses = (ecm_double *)ecm_malloc(sizeof(ecm_double) * n_alloced);

	while (ecm_fgets(buf, 256, fp)) {
		ecm_uint32	rank;
		ecm_double	prob;

		if (*buf == '#')
			continue;
		if (ecm_sscanf(buf, "%u %*u %lf", &rank, &prob) != 2) {
			ERR("%s: invalid format", buf);
			continue;
		}
		if (n_inputs >= n_alloced) {
			accesses = (ecm_double *)ecm_realloc(accesses, sizeof(ecm_double) * (n_alloced + ALLOC_CHUNK));
			n_alloced += ALLOC_CHUNK;
		}
		accesses[n_inputs++] = prob;
		if (n_inputs != rank) {
			ERR("%d: invalid rank", rank);
		}
	}

	ranks = (ecm_double *)ecm_malloc(sizeof(ecm_double) * n_inputs);
	for (i = 0; i < n_inputs; i++) {
		ranks[i] = i + 1;
	}

	return ECM_TRUE;
}

static void
do_logarithm(void)
{
	ecm_uint32	i;

	for (i = 0; i < n_inputs; i++) {
		ranks[i] = log(ranks[i]);
		accesses[i] = log(accesses[i]);
	}
}

static void
do_single_fitting(double *pc0, double *pc1, double *psumsq)
{
	ecm_double	cov00, cov01, cov11;

	gsl_fit_linear(ranks, 1, accesses, 1, n_inputs, pc0, pc1, &cov00, &cov01, &cov11, psumsq);
}

static void
do_dual_fitting(ecm_uint32 *pn_break,
		ecm_double *pc0_front, ecm_double *pc1_front,
		ecm_double *pc0_back, ecm_double *pc1_back,
		ecm_double *psumsq)
{
	ecm_double	c0_front, c1_front, c0_back, c1_back;
	ecm_double	sumsq = -1;
	ecm_double	cov00, cov01, cov11;
	ecm_uint32	i, start, end;
	ecm_uint32	n_break;

	start = n_inputs / 4;
	end = n_inputs / 4 * 3;

	for (i = start; i < end; i++) {
		ecm_double	l_c0_front, l_c1_front;
		ecm_double	l_c0_back, l_c1_back;
		ecm_double	l_sumsq_front, l_sumsq_back;

		/* front */
		gsl_fit_linear(ranks, 1, accesses, 1, i, &l_c0_front, &l_c1_front,
			       &cov00, &cov01, &cov11, &l_sumsq_front);
		gsl_fit_linear(ranks + i, 1, accesses + i, 1, n_inputs - i, &l_c0_back, &l_c1_back,
			       &cov00, &cov01, &cov11, &l_sumsq_back);

		if (sumsq < 0 || sumsq > l_sumsq_front + l_sumsq_back) {
			c0_front = l_c0_front;
			c1_front = l_c1_front;
			c0_back = l_c0_back;
			c1_back = l_c1_back;
			sumsq = l_sumsq_front + l_sumsq_back;
			n_break = i;
		}
	}

	*pc0_front = c0_front;
	*pc1_front = c1_front;
	*pc0_back = c0_back;
	*pc1_back = c1_back;
	*psumsq = sumsq;
	*pn_break = n_break;
}

static FILE *
setup_output(void)
{
	FILE	*fp;

	fp = ecm_fopen("output.txt", "w");
	if (fp == NULL) {
		ERR("cannot create output");
		return NULL;
	}

	return fp;
}

static void
report_single_fitting(ecm_double c0, ecm_double c1, ecm_double sumsq)
{
	FILE	*fp;
	ecm_uint32	i;

	if ((fp = setup_output()) == NULL)
		return;

	ecm_fprintf(fp, "# single fitting\n");
	ecm_fprintf(fp, "# c0: %.lf\n", c0);
	ecm_fprintf(fp, "# c1: %.lf\n", c1);
	ecm_fprintf(fp, "# sumsq: %.lf\n", sumsq);
	ecm_fprintf(fp, "# n_inputs: %u\n", n_inputs);

	for (i = 0; i < n_inputs; i++) {
		ecm_double	y = c0 + c1 * ranks[i];
		ecm_double	prob;

		prob = exp(y);
		ecm_fprintf(fp, "%u %8.7f\n", i + 1, prob);
	}

	ecm_fclose(fp);
}

static void
report_dual_fitting(ecm_uint32 n_break, ecm_double c0_front, ecm_double c1_front,
		    ecm_double c0_back, ecm_double c1_back, ecm_double sumsq)
{
	FILE	*fp;
	ecm_uint32	i;

	if ((fp = setup_output()) == NULL)
		return;

	ecm_fprintf(fp, "# dual fitting\n");
	ecm_fprintf(fp, "# c0_front: %lf\n", c0_front);
	ecm_fprintf(fp, "# c1_front: %lf\n", c1_front);
	ecm_fprintf(fp, "# c0_back: %lf\n", c0_back);
	ecm_fprintf(fp, "# c1_back: %lf\n", c1_back);
	ecm_fprintf(fp, "# sumsq: %.lf\n", sumsq);
	ecm_fprintf(fp, "# n_break: %u\n", n_break);
	ecm_fprintf(fp, "# n_inputs: %u\n", n_inputs);

	for (i = 0; i < n_break; i++) {
		ecm_double	y = c0_front + c1_front * ranks[i];
		ecm_double	prob;

		prob = exp(y);
		ecm_fprintf(fp, "%u %8.7f\n", i + 1, prob);
	}

	for (i = n_break; i < n_inputs; i++) {
		ecm_double	y = c0_back + c1_back * ranks[i];
		ecm_double	prob;

		prob = exp(y);
		ecm_fprintf(fp, "%u %8.7f\n", i + 1, prob);
	}

	ecm_fclose(fp);
}

static void
do_fitting(void)
{
	ecm_double	s_c0, s_c1, s_sumsq;
	ecm_double	d_c0_front, d_c1_front, d_c0_back, d_c1_back, d_sumsq;
	ecm_uint32	n_break;

	do_single_fitting(&s_c0, &s_c1, &s_sumsq);
	do_dual_fitting(&n_break, &d_c0_front, &d_c1_front, &d_c0_back, &d_c1_back, &d_sumsq);

	if (s_sumsq <= d_sumsq) {
		report_single_fitting(s_c0, s_c1, s_sumsq);
	}
	else {
		report_dual_fitting(n_break, d_c0_front, d_c1_front, d_c0_back, d_c1_back, s_sumsq);
	}
}

static void
parse_args(int argc, char *argv[])
{
	int	c;

	while ((c = ecm_getopt(argc, argv, "h")) != -1) {
		switch (c) {
		case 'h':
			usage();
			ecm_exit(0);
		default:
			usage();
			ecm_exit(1);
		}
	}

	if (ecm_optind == argc) {
		usage();
		ecm_exit(1);
	}

	inputfile = ecm_strdup(argv[ecm_optind]);
}

int
main(int argc, char *argv[])
{
	ECM_initialize();

	ecm_set_output_file(ecm_stderr);
	ecm_set_verbose_mode(ECM_TRUE);

	parse_args(argc, argv);

	if (!read_input())
		return 2;

	do_logarithm();

	do_fitting();

	return 0;
}
